import pandas as pd
import logging
import json
import os
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass

from .casualEdge import CausalityEdge
from .systemDependencyGraph import SystemDependencyGraphBuilder
from ..casual.PCMCIPlusEngine import PCMCIPlusEngine
from source.data_processor.utils.data_utils import ensure_datetime_column

# 配置日志
logger = logging.getLogger(__name__)


@dataclass
class CausalityNode:
    """因果节点数据类"""

    pod: str
    metric: str

    @property
    def id(self) -> str:
        """获取节点唯一标识符"""
        # 这里为了兼容性暂时保留使用&作为分隔符
        # 在实际存储到Neo4j时将通过generate_metric_id转换为统一格式
        return f"{self.pod}&{self.metric}"

    def __hash__(self) -> int:
        """支持将节点用作字典键或集合元素"""
        return hash(self.id)

    def __eq__(self, other) -> bool:
        """支持节点比较"""
        if not isinstance(other, CausalityNode):
            return False
        return self.id == other.id


class CausalGraph(SystemDependencyGraphBuilder):
    """
    因果图构建器，在系统依赖图的基础上添加指标间的因果关系
    该类仅处理单一阶段的因果图 (故障前或故障后)
    """

    def __init__(
        self,
        uri,
        user,
        password,
        cache_dir=None,
        batch_size=2000,
        min_data_points=1,
        phase="pre_anomaly",
        active_graph_id: str = None,
        output_dir: str = None,
    ):
        """
        初始化因果图构建器

        参数:
            uri: Neo4j数据库URI
            user: Neo4j用户名
            password: Neo4j密码
            cache_dir: 缓存目录，如果指定则启用缓存
            batch_size: 批处理大小
            min_data_points: 最小数据点数，用于因果分析
            phase: 当前图的阶段，"pre_anomaly"表示故障前，"post_anomaly"表示故障后
            active_graph_id: (可选) 此实例查询时默认操作的图ID。通常由CompositeCausalGraph设置。
            output_dir: 输出目录，用于存储JSON状态文件
        """
        # 调用父类初始化方法, 传递active_graph_id和output_dir
        super().__init__(
            uri,
            user,
            password,
            cache_dir,
            batch_size,
            active_graph_id=active_graph_id,
            output_dir=output_dir,
        )

        # 因果分析相关数据结构
        self.causal_edges: List[CausalityEdge] = []  # 存储因果关系边
        self.metrics_data_by_entity = {  # 按实体类型存储指标时序数据
            "container": {},  # container指标数据 {container_id: {metric_name: series}}
            "interface": {},  # interface指标数据 {interface_id: {metric_name: series}}
        }
        self.casual_engine = None  # 存储引擎实例

        # 因果关系属性
        self.nodes: Dict[str, CausalityNode] = {}  # node_id -> node
        # 因果分析配置
        self.min_data_points = min_data_points  # 最小数据点数
        self.max_lag = 4  # 最大时滞
        self.alpha = 0.1  # 显著性水平

        # 阶段标识
        self.phase = phase  # 当前图的阶段

    def _create_indices(self, session):
        """创建必要的索引"""
        # 首先调用父类的索引创建方法
        super()._create_indices(session)

        indices = [
            "CREATE INDEX IF NOT EXISTS FOR (s:Service) ON (s.id, s.graph_id, s.original_id)",
            "CREATE INDEX IF NOT EXISTS FOR (i:Interface) ON (i.id, i.graph_id, i.original_id)",
            "CREATE INDEX IF NOT EXISTS FOR (c:Container) ON (c.id, c.graph_id, c.original_id)",
            "CREATE INDEX IF NOT EXISTS FOR (n:Node) ON (n.id, n.graph_id, s.original_id)",
            "CREATE INDEX IF NOT EXISTS FOR (m:Metric) ON (m.id, m.graph_id, m.original_id)",
            "CREATE INDEX IF NOT EXISTS FOR ()-[r:METRIC_CAUSE]-() ON (r.method, r.strength, r.anomaly_phase, r.dependency_graph_id)",
        ]

        for index_query in indices:
            try:
                session.run(index_query)
            except Exception as e:
                logger.error(f"创建索引时发生错误: {e}")

    def _determine_entity_type(self, cmdb_id):
        """根据cmdb_id判断实体类型"""
        if pd.isna(cmdb_id):
            return None

        cmdb_id = str(cmdb_id).lower()
        if "docker" in cmdb_id or "container" in cmdb_id:
            return "container"
        elif "service" in cmdb_id or "/api/" in cmdb_id or "interface" in cmdb_id:
            return "interface"
        else:
            return "container"  # 默认为container

    def register_causal_engines(self):
        """注册因果分析引擎"""
        # 创建配置对象（简化版，实际使用时可能需要更完整的配置）
        config = type(
            "Config",
            (),
            {
                "PC_ALPHA": self.alpha,
                "MAX_LAG": self.max_lag,
                "PCMCI_TEST_TYPE": "linear",
                "MIN_DATA_POINTS": self.min_data_points,
            },
        )

        self.casual_engine = PCMCIPlusEngine(config)

    def run_causal_analysis(self, experiment_data, metric_result):
        """
        运行因果分析并处理结果

        参数:
            experiment_data: 实验数据对象
            metric_result: 指标异常检测结果

        返回:
            List[CausalityEdge]: 因果边列表
        """
        if self.casual_engine is None:
            logger.error("因果分析引擎未初始化，请先调用register_causal_engines")
            return []

        logger.info(f"开始运行{self.phase}阶段的因果分析...")

        # 获取图ID
        graph_id = experiment_data.anomaly_id

        # 运行因果分析
        causal_edges = self.casual_engine.analyze_causality(
            self, experiment_data, metric_result
        )

        # 后处理：如果多条边具有相同的源和目标节点，则仅保留p_value最小的边
        if causal_edges:
            logger.info(f"原始因果关系数量: {len(causal_edges)}")
            min_p_value_edges_map = {}
            for edge in causal_edges:
                edge_key = (
                    edge.cause_pod,
                    edge.cause_metric,
                    edge.effect_pod,
                    edge.effect_metric,
                )
                if (
                    edge_key not in min_p_value_edges_map
                    or (
                        edge.p_value is not None
                        and min_p_value_edges_map[edge_key].p_value is not None
                        and edge.p_value < min_p_value_edges_map[edge_key].p_value
                    )
                    or (
                        edge.p_value is not None
                        and min_p_value_edges_map[edge_key].p_value is None
                    )
                ):  # 当前边有p_value，旧边没有
                    min_p_value_edges_map[edge_key] = edge

            filtered_causal_edges = list(min_p_value_edges_map.values())
            if len(filtered_causal_edges) < len(causal_edges):
                logger.info(
                    f"后处理后因果关系数量: {len(filtered_causal_edges)} (移除了 {len(causal_edges) - len(filtered_causal_edges)} 条重复边)"
                )
            causal_edges = filtered_causal_edges

        # 标记阶段
        for edge in causal_edges:
            edge.anomaly_phase = self.phase

        if not causal_edges:
            logger.warning(f"{self.phase}阶段没有发现任何因果关系")
            return []

        logger.info(f"{self.phase}阶段发现 {len(causal_edges)} 个因果关系")

        # 存储获取到的边
        self.causal_edges = causal_edges

        # # 因果关系保存到Neo4j数据库的逻辑保留
        # if graph_id: # graph_id here is experiment_data.anomaly_id
        #     # self.save_causal_edges_to_cache(causal_edges, graph_id) # REMOVED - will be handled by build_causal_graph via _update_output_json
        #     self.save_causal_edges_to_neo4j(causal_edges, graph_id) # graph_id is the base anomaly_id for Neo4j relations
        # else:
        #     logger.warning("没有有效的图ID (experiment_data.anomaly_id)，无法将因果关系保存到Neo4j")

        return causal_edges

    def save_causal_edges_to_neo4j(self, edges: List[CausalityEdge], graph_id: str):
        """
        将因果关系保存到Neo4j数据库

        参数:
            edges: 因果边列表
            graph_id: 图ID
        """
        if not edges:
            logger.warning("没有因果关系可保存到Neo4j")
            return

        logger.info(f"开始将 {len(edges)} 个{self.phase}阶段因果关系保存到Neo4j")

        try:
            with self.driver.session() as session:
                # 批量添加因果关系
                for i in range(0, len(edges), self.batch_size):
                    batch = edges[i : i + self.batch_size]

                    # 构建Cypher参数
                    params = []
                    for edge in batch:
                        # 根据阶段确定依赖图ID (这个ID用于METRIC_CAUSE关系的属性)
                        dependency_graph_id_for_relation_prop = (
                            f"{graph_id}_{self.phase}"
                        )

                        # Metric节点在创建时，其ID内嵌的graph_id是带phase的
                        # 因此，这里生成cause_id和effect_id时，必须使用带phase的graph_id
                        metric_node_creation_graph_id = f"{graph_id}_{self.phase}"

                        cause_id = self.generate_metric_id(
                            edge.cause_pod,
                            edge.cause_metric,
                            metric_node_creation_graph_id,
                        )
                        effect_id = self.generate_metric_id(
                            edge.effect_pod,
                            edge.effect_metric,
                            metric_node_creation_graph_id,
                        )

                        params.append(
                            {
                                "cause_id": cause_id,
                                "effect_id": effect_id,
                                "method": edge.method,
                                "strength": edge.strength,
                                "time_lag": edge.time_lag,
                                "p_value": edge.p_value,
                                "graph_id": graph_id,  # 这个graph_id是基础的anomaly_id，可能用于关系自身的标识
                                "dependency_graph_id": dependency_graph_id_for_relation_prop,  # 关系属性，指向带phase的图
                                "from_entity_type": edge.from_entity_type,
                                "to_entity_type": edge.to_entity_type,
                                "anomaly_phase": self.phase,
                            }
                        )

                    # 执行Cypher查询
                    query = """
                    UNWIND $params AS param
                    MATCH (a:Metric {id: param.cause_id}), 
                          (b:Metric {id: param.effect_id})
                    MERGE (a)-[r:METRIC_CAUSE {
                        method: param.method,
                        strength: param.strength,
                        time_lag: param.time_lag,
                        p_value: param.p_value,
                        graph_id: param.graph_id,
                        dependency_graph_id: param.dependency_graph_id,
                        from_entity_type: param.from_entity_type,
                        to_entity_type: param.to_entity_type,
                        anomaly_phase: param.anomaly_phase  
                    }]->(b)
                    """

                    session.run(query, {"params": params})

                logger.info(
                    f"成功将 {len(edges)} 个{self.phase}阶段因果关系保存到Neo4j"
                )

        except Exception as e:
            logger.error(f"保存因果关系到Neo4j时出错: {str(e)}")
            import traceback

            logger.error(traceback.format_exc())

    def build_causal_graph(
        self, experiment_data, metric_result=None, trace_result=None, time_range=None
    ):
        """
        构建单一阶段的因果图, 集成新的缓存和JSON状态更新逻辑。

        参数:
            experiment_data: 实验数据对象
            metric_result: 指标异常检测结果，用于筛选异常指标
            trace_result: 链路异常检测结果（已不再使用）
            time_range: 时间范围元组 (start_time, end_time)，用于过滤数据

        返回:
            bool: 构建成功返回True，否则返回False
        """
        # 初始化参数
        graph_id = experiment_data.anomaly_id  # Base anomaly_id
        phase_graph_id = f"{graph_id}_{self.phase}"  # Graph ID for this specific phase
        self.active_graph_id = phase_graph_id  # 设置活动图ID

        logger.info(f"开始构建{self.phase}阶段的因果图, graph_id={phase_graph_id}")
        if time_range:
            logger.info(f"使用时间范围: {time_range[0]} 到 {time_range[1]}")

        # 尝试从缓存或JSON文件恢复已构建的图
        if self._try_restore_from_cache_or_json(phase_graph_id):
            return True

        # 1. 构建依赖图
        dependency_success = self._build_dependency_graph(
            experiment_data, metric_result, trace_result, phase_graph_id, time_range
        )
        if not dependency_success:
            return False

        # 2. 获取适用于因果分析的数据（过滤指标数据）
        analysis_data = self._prepare_data_for_causal_analysis(
            experiment_data, time_range
        )
        if not self._check_data_sufficiency(analysis_data):
            # 如果数据不足，依赖图已构建，返回成功
            return True

        # 3. 执行因果分析
        self._perform_causal_analysis(analysis_data, metric_result)

        # 4. 更新JSON状态
        self._update_json_state(phase_graph_id)

        logger.info(f"{self.phase}阶段因果图构建流程完成, graph_id={phase_graph_id}")
        return True

    def _try_restore_from_cache_or_json(self, phase_graph_id):
        """
        尝试从缓存或JSON状态文件恢复已构建的图

        参数:
            phase_graph_id: 带阶段标识的图ID

        返回:
            bool: 如果成功恢复则返回True，否则返回False
        """
        # 如果未配置输出目录，无法恢复
        if not self.output_dir:
            logger.warning(
                f"output_dir 未配置 for CausalGraph。将不使用JSON状态文件进行恢复或更新。"
            )
            return False

        # 尝试读取JSON状态文件
        output_graph_dir = os.path.join(self.output_dir, "graph")
        os.makedirs(output_graph_dir, exist_ok=True)  # 确保输出目录存在
        output_json_path = os.path.join(output_graph_dir, f"{phase_graph_id}.json")

        if not os.path.exists(output_json_path):
            return False

        try:
            with open(output_json_path, "r") as f:
                status_data = json.load(f)
            logger.info(f"发现已存在的Output JSON状态文件: {output_json_path}")

            # 检查构建阶段
            build_phase = status_data.get("build_phase")
            pkl_cache_file_path = status_data.get("pkl_cache_file_path")

            # 如果已完成因果图构建且缓存文件存在
            if (
                build_phase == "causal_graph_built"
                and pkl_cache_file_path
                and os.path.exists(pkl_cache_file_path)
            ):

                logger.info(
                    f"检测到构建阶段 '{build_phase}'，尝试从缓存加载依赖图和因果边..."
                )

                # 加载依赖图
                if self.load_graph_from_cache(phase_graph_id):
                    # 加载因果边
                    causal_graph_data = status_data.get("causal_graph_data", {})
                    loaded_edges_dicts = causal_graph_data.get("causal_edges", [])
                    self.causal_edges = [
                        CausalityEdge(**edge_dict) for edge_dict in loaded_edges_dicts
                    ]

                    logger.info(
                        f"{self.phase}阶段因果图从缓存和JSON状态完全恢复, ({len(self.causal_edges)} 因果边)"
                    )
                    return True
                else:
                    logger.warning(
                        f"无法从PKL ({pkl_cache_file_path}) 加载依赖结构，将继续执行构建流程。"
                    )
        except Exception as e:
            logger.error(
                f"读取或解析Output JSON ({output_json_path}) 时发生错误: {e}。将继续执行构建流程。"
            )

        return False

    def _build_dependency_graph(
        self, experiment_data, metric_result, trace_result, phase_graph_id, time_range
    ):
        """
        构建系统依赖图

        参数:
            experiment_data: 实验数据对象
            metric_result: 指标异常检测结果
            trace_result: 链路异常检测结果
            phase_graph_id: 带阶段标识的图ID
            time_range: 时间范围元组

        返回:
            bool: 构建成功返回True，否则返回False
        """
        logger.info(f"构建{self.phase}阶段的依赖图部分...")

        # 调用父类方法构建依赖图
        dependency_success, self.pkl_path_from_dependency_build = super().build_graph(
            experiment_data,
            metric_result,
            custom_graph_id=phase_graph_id,
            time_range=time_range,
        )

        if not dependency_success:
            logger.error(f"构建{self.phase}阶段系统依赖图失败")
            return False

        logger.info(
            f"已成功构建{self.phase}阶段系统依赖图 (PKL: {self.pkl_path_from_dependency_build})"
        )
        return True

    def _prepare_data_for_causal_analysis(self, experiment_data, time_range):
        """
        准备用于因果分析的数据（过滤指标数据）

        参数:
            experiment_data: 实验数据对象
            time_range: 时间范围元组

        返回:
            ExperimentData: 过滤后的实验数据对象
        """
        # 如果提供了时间范围，过滤数据
        if time_range:
            # 使用ExperimentData的过滤方法创建新的对象
            filtered_data = experiment_data.filter_by_time_range(time_range)
            logger.info(
                f"为因果分析准备了过滤后的数据: {len(filtered_data.metrics_df)}/{len(experiment_data.metrics_df)} 行指标数据"
            )
            return filtered_data

        # 如果没有时间范围，直接使用原始数据
        return experiment_data

    def _check_data_sufficiency(self, analysis_data):
        """
        检查数据是否足够进行因果分析

        参数:
            analysis_data: 实验数据对象

        返回:
            bool: 如果数据足够则返回True，否则返回False
        """
        if (
            analysis_data.metrics_df.empty
            or len(analysis_data.metrics_df) < self.min_data_points
        ):
            logger.warning(
                f"{self.phase}阶段指标数据不足 (行数: {len(analysis_data.metrics_df)})，跳过因果分析"
            )
            return False
        return True

    def _perform_causal_analysis(self, analysis_data, metric_result):
        """
        执行因果分析

        参数:
            analysis_data: 实验数据对象
            metric_result: 指标异常检测结果

        返回:
            bool: 分析成功返回True，否则返回False
        """
        # 初始化因果分析引擎
        self.register_causal_engines()
        logger.info(f"运行{self.phase}阶段的因果分析...")

        # 执行因果分析
        self.run_causal_analysis(analysis_data, metric_result)

        # 记录结果
        if not self.causal_edges:
            logger.warning(f"{self.phase}阶段没有发现因果关系")
            return False
        else:
            logger.info(
                f"{self.phase}阶段因果分析完成, 发现 {len(self.causal_edges)} 条因果关系"
            )
            return True

    def _update_json_state(self, phase_graph_id):
        """
        更新JSON状态文件

        参数:
            phase_graph_id: 带阶段标识的图ID
        """
        if not self.output_dir:
            return

        # 更新JSON状态
        if self.pkl_path_from_dependency_build:
            logger.info(
                f"更新JSON状态为 causal_graph_built, pkl_path: {self.pkl_path_from_dependency_build}"
            )
            self._update_output_json(
                phase_graph_id,
                "causal_graph_built",
                self.pkl_path_from_dependency_build,
                causal_data={"causal_edges": self.causal_edges},
            )
        else:
            logger.warning(f"由于依赖图构建阶段未提供有效的PKL路径，未能更新JSON状态")
